import mindspore.nn as nn
import mindspore.ops as ops
import mindspore as ms
from .ijepa_3d import IJEPA3DConfig, IJEPA3DForPreTraining
from .vit_3d import ViT3DConfig, ViT3DModel
from .mae_3d import ViTMAE3DConfig, ViTMAE3DModel, ViTMAE3DForPreTraining
from .cls_haed import ClassificationHead
from .clip import CLIPModel, ClipConfig, TextConfig
from .mae_mm import ViTMAEConfigMM, ViTMAEForPreTrainingMM

def build_model(args):
    vit_args = dict(
        image_size=tuple(args.img_size), 
        patch_size=tuple(args.patch_size),
        hidden_size=args.embed_dim,
        intermediate_size=args.embed_dim * 4,
        num_hidden_layers=args.depth,
        num_attention_heads=args.num_heads,
        sin_cos_embed=args.sin_cos_embed,
        recompute=args.recompute
    )
    if args.exp_type == 'mae':
        cfg = ViTMAE3DConfig(
            mask_ratio=args.mask_ratio,
            decoder_hidden_size=args.decoder_embed_dim,
            decoder_intermediate_size=args.decoder_embed_dim * 4,
            decoder_num_hidden_layers=args.decoder_depth,
            decoder_num_attention_heads=args.decoder_num_heads,
            **vit_args
        )
        model = ViTMAE3DForPreTraining(cfg)
    elif args.exp_type == 'ijepa':
        cfg = IJEPA3DConfig(
            M=args.M,
            decoder_num_hidden_layers=args.decoder_depth,
            **vit_args
        )
        model = IJEPA3DForPreTraining(cfg)
    elif args.exp_type == 'clip':
        vision_config = ViT3DConfig(
                global_pool=False,
                param_init_type=ms.float16,
                **vit_args
                )
        text_config = TextConfig(
            bert_path=args.bert_path,
            max_length=args.max_length
        )
        clip_config = ClipConfig(
            text_config=text_config,
            vision_config=vision_config,
            projection_dim=args.projection_dim
        )
        model = CLIPModel(clip_config)
    elif args.exp_type == 'cls':
        vit_config = ViT3DConfig(
            use_mean_pooling=True,
            **vit_args
            )
        model = ClassificationHead(ViT3DModel(vit_config), args.embed_dim, num_classes=10)
    elif args.exp_type == 'mm':
        cfg = ViTMAEConfigMM(
            mask_ratio=args.mask_ratio,
            modality_list=tuple(args.modality_list),
            num_channels=args.total_dim,
            seq_len=args.seq_len,
            decoder_hidden_size=args.decoder_embed_dim,
            decoder_intermediate_size=args.decoder_embed_dim * 4,
            decoder_num_hidden_layers=args.decoder_depth,
            decoder_num_attention_heads=args.decoder_num_heads,
            norm_pix_loss=args.norm_pix_loss,
            loss_type=args.loss_type,
            **vit_args
        )
        model = ViTMAEForPreTrainingMM(cfg)
    else:
        raise NotImplementedError

    return model